import os
from utils.detection_3d.visualization import parse_npy
from utils.detection_3d.calculate_utils import *
import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
# import openpyxl


def draw_pic_for_loop(img_data, target_pic=None):
    plt.cla()
    for i in img_data:
        plt.plot(i)
    if not target_pic is None:
        plt.savefig(target_pic)
    plt.show()


def draw_pic_for_loop_with_name(img_data, target_pic=None):
    plt.cla()
    for i in img_data:
        plt.plot(i[1])
    if not target_pic is None:
        plt.savefig(target_pic)
    plt.show()


def analyse_shape(z_point, real_bottom, area=None, draw_file=None):
    pin_shape = []
    data = np.ones(z_point.shape)
    for i in np.arange(z_point.max(), real_bottom, -0.0125):
        index = z_point >= i
        data[index] = 1
        data[~index] = 0
        pin_shape.append(np.sum(data))
    if not (draw_file is None or area is None):
        path, fin = os.path.split(draw_file)
        fname, fsuffix = os.path.splitext(fin)
        plt.cla()
        plt.plot(pin_shape)
        plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\nok", area + "_" + fname + ".png"))
    else:
        plt.cla()
        plt.plot(pin_shape)
        plt.show()


# def analyse_together(z_point, real_bottom, pinok, pinnok, filename, area):
#     pin_shape = []
#     data = np.ones(z_point.shape)
#     for i in np.arange(z_point.max(), real_bottom - 3, -0.0125):
#         index = z_point >= i
#         data[index] = 1
#         data[~index] = 0
#         pin_shape.append(np.sum(data))
#     target = "D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\label.xlsx"
#     data = openpyxl.load_workbook(target)
#     stn = data.sheetnames[0]
#     # sheetname = data.get_sheet_names
#     table = data.get_sheet_by_name(stn)
#     table = data.active
#     nrows = table.max_row
#     flag = True
#     for i in range(1, nrows + 1):
#         if table.cell(i, 1).value == filename + ".npy":
#             flag = False
#             if table.cell(i, 1 + area).value == "n":
#                 pinnok.append(pin_shape)
#             else:
#                 pinok.append(pin_shape)
#             break
#     if flag:
#         pinok.append(pin_shape)
#     return pinok, pinnok


# def test_func(npy, pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p):
#     path, fin = os.path.split(npy)
#     fname, fsuffix = os.path.splitext(fin)
#     xyz_points = parse_npy(npy)  # 得到x，y，z坐标
#     x_points, y_points, z_points = xyz_points
#
#     # 未处理的数据显示
#
#     bottom_down, bottom_up, bottom_left = get_bottom_value(x_points, y_points, z_points)
#     Z_BOTTOM = (bottom_down + bottom_up + bottom_left) / 3  # PCD面的大致底面高度
#     # ==========================AREA_1 最下面 第一个焊点的正面==============================
#     # 根据区域，大致定位，过滤数据
#
#     _x1, _y1, _z1 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 10 * 4 + 1,
#                                         x_max=x_points.max() / 10 * 6 - 0, y_min=y_points.min() + 1,
#                                         y_max=y_points.max() / 5 + 1.5)
#     # o3d.visualization.draw_geometries([numpy2pcd(_x1,_y1,_z1)], window_name='AREA_1_TEP')
#     # 根据底面高度，准确定位
#     # draw_z_points(_x1, _z1)
#     real_bottom = get_real_bottom(Z_BOTTOM, _z1)
#     x1, y1, z1 = data_filter_by_Z_value(_x1, _y1, _z1, real_bottom)
#
#     # analyse_shape(z1, 0, "pin1", npy)
#     # pin1_p, pin1_n = analyse_together(z1, 0, pin1_p, pin1_n, fname, 1)
#     pin1_p, pin1_n = analyse_together(_z1, 0, pin1_p, pin1_n, fname, 1)
#     # draw_z_points(x1, z1)
#     # o3d.visualization.draw_geometries([numpy2pcd(x1, y1, z1)], window_name='AREA_1')
#     v1_front = np.sum(z1) * CONSTANT_UP
#     print("the AREA_1 front volume is", v1_front)
#
#     # ==========================AREA_2 第二个焊点的正面==============================
#     # 根据区域，大致定位，过滤数据
#     _x2, _y2, _z2 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 2 - 13,
#                                         x_max=x_points.max() / 2 - 8, y_min=y_points.max() / 2 - 2,
#                                         y_max=y_points.max() / 20 * 11 + 0.5)
#     # o3d.visualization.draw_geometries([numpy2pcd(_x2,_y2,_z2)], window_name='AREA_2_TEP')
#     # 根据底面高度，准确定位
#     # draw_z_points(_x2, _z2)
#     real_bottom = get_real_bottom(Z_BOTTOM, _z2)
#     x2, y2, z2 = data_filter_by_Z_value(_x2, _y2, _z2, real_bottom)
#     # analyse_shape(z2, 0, "pin2", npy)
#     # pin2_p, pin2_n = analyse_together(z2, 0, pin2_p, pin2_n, fname, 2)
#     pin2_p, pin2_n = analyse_together(_z2, 0, pin2_p, pin2_n, fname, 2)
#     area_2_pcd = numpy2pcd(x2, y2, z2)
#     # draw_z_points(x2, z2)
#     # o3d.visualization.draw_geometries([area_2_pcd], window_name='AREA_2')
#     _, _, z2 = pcd2numpy(area_2_pcd)
#     v2_front = np.sum(z2) * CONSTANT_UP
#     print("the AREA_2 front volume is", v2_front)
#
#     # ==========================AREA_3 第三个焊点的正面==============================
#     # 根据区域，大致定位，过滤数据
#     _x3, _y3, _z3 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 8 - 0.5,
#                                         x_max=x_points.max() / 8 * 3 - 3.5, y_min=y_points.max() / 5 * 4 - 0.5,
#                                         y_max=y_points.max() - 0.8 + 10)
#     # o3d.visualization.draw_geometries([numpy2pcd(_x3,_y3,_z3)], window_name='AREA_3_TEP')
#     # 根据底面高度，准确定位
#     # draw_z_points(_x3, _z3)
#     real_bottom = get_real_bottom(Z_BOTTOM, _z3)
#     x3, y3, z3 = data_filter_by_Z_value(_x3, _y3, _z3, real_bottom)
#     # analyse_shape(z3, 0, "pin3", npy)
#     # pin3_p, pin3_n = analyse_together(z3, 0, pin3_p, pin3_n, fname, 3)
#     pin3_p, pin3_n = analyse_together(_z3, 0, pin3_p, pin3_n, fname, 3)
#     # draw_z_points(x3, z3)
#     area_3_pcd = numpy2pcd(x3, y3, z3)
#     # o3d.visualization.draw_geometries([area_3_pcd], window_name='AREA_3')
#     _, _, z3_front_res = pcd2numpy(area_3_pcd)
#     v3_front = np.sum(z3_front_res) * CONSTANT_UP
#     print("the AREA_3 front volume is", v3_front)
#     return pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p


# def write_name(root):
#     files = os.listdir(root)
#     target = "D:\\Programs\\data\\dataset\\filename.xlsx"
#     data = openpyxl.load_workbook(target)
#     stn = data.sheetnames[0]
#     # sheetname = data.get_sheet_names
#     table = data.get_sheet_by_name(stn)
#     table = data.active
#     nrows = table.max_row
#     for i in files:
#         if i[-3:] == "npy" and i.count("_up_") > 0:
#             table.cell(nrows + 1, 1).value = i
#             nrows = nrows + 1
#     data.save(target)


# def test_batch(root):
#     pin1_n = []
#     pin1_p = []
#     pin2_n = []
#     pin2_p = []
#     pin3_n = []
#     pin3_p = []
#     file = os.listdir(root)
#     for i in file:
#         try:
#             if i.count("_up_") > 0 and i[-3:] == "npy":
#                 print(os.path.join(root, i))
#                 pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_func(os.path.join(root, i), pin1_n, pin1_p,
#                                                                            pin2_n, pin2_p, pin3_n, pin3_p)
#         except Exception as e:
#             print(os.path.join(root, i), e)
#     return pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p


# def analyse_together_return_with_name(z_point, real_bottom, pinok, pinnok, filename, area):
#     pin_shape = []
#     data = np.ones(z_point.shape)
#     for i in np.arange(z_point.max(), real_bottom - 3, -0.0125):
#         index = z_point >= i
#         data[index] = 1
#         data[~index] = 0
#         pin_shape.append(np.sum(data))
#     target = "D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\label.xlsx"
#     data = openpyxl.load_workbook(target)
#     stn = data.sheetnames[0]
#     # sheetname = data.get_sheet_names
#     table = data.get_sheet_by_name(stn)
#     table = data.active
#     nrows = table.max_row
#     flag = True
#     for i in range(1, nrows + 1):
#         if table.cell(i, 1).value == filename + ".npy":
#             flag = False
#             if table.cell(i, 1 + area).value == "n":
#                 pinnok.append((filename, pin_shape))
#             else:
#                 pinok.append((filename, pin_shape))
#             break
#     if flag:
#         pinok.append((filename, pin_shape))
#     return pinok, pinnok


# def test_func_return_with_name(npy, pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p):
#     path, fin = os.path.split(npy)
#     fname, fsuffix = os.path.splitext(fin)
#     xyz_points = parse_npy(npy)  # 得到x，y，z坐标
#     x_points, y_points, z_points = xyz_points
#
#     # 未处理的数据显示
#
#     bottom_down, bottom_up, bottom_left = get_bottom_value(x_points, y_points, z_points)
#     Z_BOTTOM = (bottom_down + bottom_up + bottom_left) / 3  # PCD面的大致底面高度
#     # ==========================AREA_1 最下面 第一个焊点的正面==============================
#     # 根据区域，大致定位，过滤数据
#
#     _x1, _y1, _z1 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 10 * 4 + 1,
#                                         x_max=x_points.max() / 10 * 6 - 0, y_min=y_points.min() + 1,
#                                         y_max=y_points.max() / 5 + 1.5)
#     # o3d.visualization.draw_geometries([numpy2pcd(_x1,_y1,_z1)], window_name='AREA_1_TEP')
#     # 根据底面高度，准确定位
#     # draw_z_points(_x1, _z1)
#     real_bottom = get_real_bottom(Z_BOTTOM, _z1)
#     x1, y1, z1 = data_filter_by_Z_value(_x1, _y1, _z1, real_bottom)
#
#     # analyse_shape(z1, 0, "pin1", npy)
#     # pin1_p, pin1_n = analyse_together_return_with_name(z1, 0, pin1_p, pin1_n, fname, 1)
#     pin1_p, pin1_n = analyse_together_return_with_name(_z1, 0, pin1_p, pin1_n, fname, 1)
#     # draw_z_points(x1, z1)
#     # o3d.visualization.draw_geometries([numpy2pcd(x1, y1, z1)], window_name='AREA_1')
#     v1_front = np.sum(z1) * CONSTANT_UP
#     print("the AREA_1 front volume is", v1_front)
#
#     # ==========================AREA_2 第二个焊点的正面==============================
#     # 根据区域，大致定位，过滤数据
#     _x2, _y2, _z2 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 2 - 13,
#                                         x_max=x_points.max() / 2 - 8, y_min=y_points.max() / 2 - 2,
#                                         y_max=y_points.max() / 20 * 11 + 0.5)
#     # o3d.visualization.draw_geometries([numpy2pcd(_x2,_y2,_z2)], window_name='AREA_2_TEP')
#     # 根据底面高度，准确定位
#     # draw_z_points(_x2, _z2)
#     real_bottom = get_real_bottom(Z_BOTTOM, _z2)
#     x2, y2, z2 = data_filter_by_Z_value(_x2, _y2, _z2, real_bottom)
#     # analyse_shape(z2, 0, "pin2", npy)
#     # pin2_p, pin2_n = analyse_together_return_with_name(z2, 0, pin2_p, pin2_n, fname, 2)
#     pin2_p, pin2_n = analyse_together_return_with_name(_z2, 0, pin2_p, pin2_n, fname, 2)
#     area_2_pcd = numpy2pcd(x2, y2, z2)
#     # draw_z_points(x2, z2)
#     # o3d.visualization.draw_geometries([area_2_pcd], window_name='AREA_2')
#     _, _, z2 = pcd2numpy(area_2_pcd)
#     v2_front = np.sum(z2) * CONSTANT_UP
#     print("the AREA_2 front volume is", v2_front)
#
#     # ==========================AREA_3 第三个焊点的正面==============================
#     # 根据区域，大致定位，过滤数据
#     _x3, _y3, _z3 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 8 - 0.5,
#                                         x_max=x_points.max() / 8 * 3 - 3.5, y_min=y_points.max() / 5 * 4 - 0.5,
#                                         y_max=y_points.max() - 0.8 + 10)
#     # o3d.visualization.draw_geometries([numpy2pcd(_x3,_y3,_z3)], window_name='AREA_3_TEP')
#     # 根据底面高度，准确定位
#     # draw_z_points(_x3, _z3)
#     real_bottom = get_real_bottom(Z_BOTTOM, _z3)
#     x3, y3, z3 = data_filter_by_Z_value(_x3, _y3, _z3, real_bottom)
#     # analyse_shape(z3, 0, "pin3", npy)
#     # pin3_p, pin3_n = analyse_together_return_with_name(z3, 0, pin3_p, pin3_n, fname, 3)
#     pin3_p, pin3_n = analyse_together_return_with_name(_z3, 0, pin3_p, pin3_n, fname, 3)
#     # draw_z_points(x3, z3)
#     area_3_pcd = numpy2pcd(x3, y3, z3)
#     # o3d.visualization.draw_geometries([area_3_pcd], window_name='AREA_3')
#     _, _, z3_front_res = pcd2numpy(area_3_pcd)
#     v3_front = np.sum(z3_front_res) * CONSTANT_UP
#     print("the AREA_3 front volume is", v3_front)
#     return pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p


# def test_batch_return_with_name(root):
#     pin1_n = []
#     pin1_p = []
#     pin2_n = []
#     pin2_p = []
#     pin3_n = []
#     pin3_p = []
#     file = os.listdir(root)
#     for i in file:
#         try:
#             if i.count("_up_") > 0 and i[-3:] == "npy":
#                 print(os.path.join(root, i))
#                 pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_func_return_with_name(os.path.join(root, i),
#                                                                                             pin1_n, pin1_p, pin2_n,
#                                                                                             pin2_p, pin3_n, pin3_p)
#         except Exception as e:
#             print(os.path.join(root, i), e)
#     return pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p


# if __name__ == "__main__":
#     import os
#     from visualization import parse_npy
#     from calculate_utils import *
#     import open3d as o3d
#     import numpy as np
#
#     # 3505277134ML085B085301118062200861_up_2022-08-09-10-49-52
#     # write_name(root)
#     root = "D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\data"
#     # write_name(root)
#     pin1_n = []
#     pin1_p = []
#     pin2_n = []
#     pin2_p = []
#     pin3_n = []
#     pin3_p = []
#     pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_batch(
#         "D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\pin_nok\\pin1")
#     plt.cla()
#     for i in pin1_n:
#         plt.plot(i)
#     plt.show()
#     pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_batch(
#         "D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\pin_nok\\pin2")
#     plt.cla()
#     for i in pin2_n:
#         plt.plot(i)
#     plt.show()
#     pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_batch(
#         "D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\pin_nok\\pin3")
#     plt.cla()
#     for i in pin3_n:
#         plt.plot(i)
#     plt.show()
#     # file = os.listdir(root)
#     # for i in file:
#     #     try:
#     #         if i.count("_up_") > 0 and i[-3:] == "npy":
#     #             print(os.path.join(root, i))
#     #             pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_func(os.path.join(root, i), pin1_n, pin1_p,
#     #                                                                        pin2_n, pin2_p, pin3_n, pin3_p)
#     #     except Exception as e:
#     #         print(os.path.join(root, i), e)
#     # root = "D:\\Programs\\data\\dataset\\20220902"
#     # file = os.listdir(root)
#     # for i in file:
#     #     try:
#     #         if i.count("_up_") > 0 and i[-3:] == "npy":
#     #             print(os.path.join(root, i))
#     #             pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_func(os.path.join(root, i), pin1_n, pin1_p,
#     #                                                                        pin2_n, pin2_p, pin3_n, pin3_p)
#     #     except Exception as e:
#     #         print(os.path.join(root, i), e)
#     # 3505277134ML085B085301118062201007_up_2022-08-09-09-01-58.npy
#     # 3505277134ML085B085301118062200973_up_2022-08-09-07-23-24.npy
#     # pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_func("D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\data\\3505277134ML085B085301118062201007_up_2022-08-09-09-01-58.npy", pin1_n, pin1_p,
#     #                                                            pin2_n, pin2_p, pin3_n, pin3_p)
#     # pin1_n, pin1_p, pin2_n, pin2_p, pin3_n, pin3_p = test_func("D:\\Programs\\data\\dataset\\20220808+20220706_fly_tin\\data\\3505277134ML085B085301118062200973_up_2022-08-09-07-23-24.npy", pin1_n, pin1_p,
#     #                                                            pin2_n, pin2_p, pin3_n, pin3_p)
#     #
#     # print(1)
#     # plt.cla()
#     # for i in pin2_n:
#     #     plt.plot(i)
#     # for i in pin2_p:
#     #     plt.plot(i)
#     # plt.show()
#     # plt.cla()
#     # for i in pin1_p:
#     #     plt.plot(i, color="green")
#     # for i in pin1_n:
#     #     plt.plot(i, color="red")
#     # # plt.show()
#     # plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\total", "pin1_1.png"))
#     # plt.cla()
#     # for i in pin2_p:
#     #     plt.plot(i, color="green")
#     # for i in pin2_n:
#     #     plt.plot(i, color="red")
#     # # plt.show()
#     # plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\total", "pin2_1.png"))
#     # plt.cla()
#     # for i in pin3_p:
#     #     plt.plot(i, color="green")
#     # for i in pin3_n:
#     #     plt.plot(i, color="red")
#     # # plt.show()
#     # plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\total", "pin3_1.png"))
#     # plt.cla()
#     # for i in pin1_n:
#     #     plt.plot(i, color="red")
#     # for i in pin1_p:
#     #     plt.plot(i, color="green")
#     # # plt.show()
#     # plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\total", "pin1_2.png"))
#     # plt.cla()
#     # for i in pin2_n:
#     #     plt.plot(i, color="red")
#     # for i in pin2_p:
#     #     plt.plot(i, color="green")
#     # # plt.show()
#     # plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\total", "pin2_2.png"))
#     # plt.cla()
#     # for i in pin3_n:
#     #     plt.plot(i, color="red")
#     # for i in pin3_p:
#     #     plt.plot(i, color="green")
#     # # plt.show()
#     # plt.savefig(os.path.join("D:\\Programs\\data\\exp\\正面pin针形态_切片\\total", "pin3_2.png"))
